import torch
import torch.nn as nn
import torch.nn.functional as F


def normal_init(module, mean=0, std=1, bias=0):
    if hasattr(module, 'weight') and module.weight is not None:
        nn.init.normal_(module.weight, mean, std)
    if hasattr(module, 'bias') and module.bias is not None:
        nn.init.constant_(module.bias, bias)


def constant_init(module, val, bias=0):
    if hasattr(module, 'weight') and module.weight is not None:
        nn.init.constant_(module.weight, val)
    if hasattr(module, 'bias') and module.bias is not None:
        nn.init.constant_(module.bias, bias)


class DySample(nn.Module):
    def __init__(self, in_channels, scale=2, style='lp', groups=4, dyscope=False):
        super().__init__()
        self.scale = scale
        self.style = style
        self.groups = groups
        assert style in ['lp', 'pl']
        if style == 'pl':
            assert in_channels >= scale ** 2 and in_channels % scale ** 2 == 0
        assert in_channels >= groups and in_channels % groups == 0

        if style == 'pl':
            in_channels = in_channels // scale ** 2
            out_channels = 2 * groups
        else:
            out_channels = 2 * groups * scale ** 2

        self.offset = nn.Conv2d(in_channels, out_channels, 1)
        normal_init(self.offset, std=0.001)
        if dyscope:
            self.scope = nn.Conv2d(in_channels, out_channels, 1)
            constant_init(self.scope, val=0.)

        self.register_buffer('init_pos', self._init_pos())

    def _init_pos(self):
        h = torch.arange((-self.scale + 1) / 2, (self.scale - 1) / 2 + 1) / self.scale
        return torch.stack(torch.meshgrid([h, h])).transpose(1, 2).repeat(1, self.groups, 1).reshape(1, -1, 1, 1)

    # 对输入图像x进行空间采样，基于给定的偏移量offset生成新的坐标，
    # 并利用这些坐标通过双线性插值从原图中取出像素值，实现图像的上采样或空间变换。
    def sample(self, x, offset):
        # 获取形状信息：首先获取offset的形状信息，即批大小B、通道数（这里忽略，用_占位）、高度H、宽度W。
        B, _, H, W = offset.shape
        # 重塑偏移量：将offset调整形状为(B, 2, -1, H, W)，意味着每个像素有两个偏移值，分别对应水平和垂直方向。
        offset = offset.view(B, 2, -1, H, W)
        # 创建网格坐标：创建一个HxW的网格，每个像素点对应一个坐标，坐标范围在[-0.5, 0.5]。
        coords_h = torch.arange(H) + 0.5
        coords_w = torch.arange(W) + 0.5
        coords = torch.stack(torch.meshgrid([coords_w, coords_h])
                             ).transpose(1, 2).unsqueeze(1).unsqueeze(0).type(x.dtype).to(x.device)
        # 归一化坐标：创建一个归一化因子(W, H)，确保坐标转换后在[-1, 1]
        # 范围内，这是双线性插值函数要求的坐标格式。
        normalizer = torch.tensor([W, H], dtype=x.dtype, device=x.device).view(1, 2, 1, 1, 1)
        coords = 2 * (coords + offset) / normalizer - 1
        # 使用F.pixel_shuffle函数对归一化后的坐标进行上采样，使其与目标放大尺度一致。
        # 接着调整坐标维度顺序并展平，以便于后续的网格采样。
        coords = F.pixel_shuffle(coords.view(B, -1, H, W), self.scale).view(
            B, 2, -1, self.scale * H, self.scale * W).permute(0, 2, 3, 4, 1).contiguous().flatten(0, 1)
        # 网格采样：将输入图像x调整形状以匹配组数和新尺寸，然后使用F.grid_sample函数是pytorch中有双线性插值的方法，
        # 根据计算出的坐标网格对图像进行双线性插值采样。这里设置了align_corners = False来控制边缘处理方式，
        # 以及padding_mode = "border"来处理超出边界的查询。
        # 最后调用view函数将结果调整形状为(B, -1, H, W)，并返回。
        # 其中-1代表自动计算该维度的大小，完成整个空间采样的过程。
        return F.grid_sample(x.reshape(B * self.groups, -1, self.scale * H, self.scale * W), coords, mode='bilinear',
                             align_corners=False, padding_mode="border").view(B, -1, self.scale * H, self.scale * W)

    def forward_lp(self, x):
        if hasattr(self, 'scope'):
            offset = self.offset(x) * self.scope(x).sigmoid() * 0.5 + self.init_pos
        else:
            offset = self.offset(x) * 0.25 + self.init_pos
        return self.sample(x, offset)

    def forward_pl(self, x):
        x_ = F.pixel_shuffle(x, self.scale)
        if hasattr(self, 'scope'):
            offset = F.pixel_unshuffle(self.offset(x_) * self.scope(x_).sigmoid(), self.scale) * 0.5 + self.init_pos
        else:
            offset = F.pixel_unshuffle(self.offset(x_), self.scale) * 0.25 + self.init_pos
        return self.sample(x, offset)

    def forward(self, x):
        if self.style == 'pl':
            return self.forward_pl(x)
        return self.forward_lp(x)
